package com.gateway.router.ratelimiter.resolver;

import com.google.common.collect.Maps;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.PrincipalNameKeyResolver;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.Map;
import java.util.Objects;

/**
 * @program: api-gateway
 * @description:
 * @author: YuKai Fan
 * @create: 2025/2/21 17:01
 **/
public class KeyResolverFactory implements KeyResolver {
    public static final String BEAN_NAME = "keyResolverMap";
    private static final String EMPTY_KEY = "____EMPTY_KEY__";

    private static final Map<String, KeyResolver> keyResolverMap;

    static {
        keyResolverMap = Maps.newHashMap();
        keyResolverMap.put(RemoteAddrKeyResolver.BEAN_NAME, new RemoteAddrKeyResolver());
        keyResolverMap.put(WholeKeyResolver.BEAN_NAME, new WholeKeyResolver());
        keyResolverMap.put(PrincipalNameKeyResolver.BEAN_NAME, new PrincipalNameKeyResolver());
    }

    @Override
    public Mono<String> resolve(ServerWebExchange exchange) {
        Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
        if (Objects.isNull(route)) {
            return Mono.just(EMPTY_KEY);
        }
        Map<String, Object> metadata = route.getMetadata();
        String keyResolver = (String) metadata.getOrDefault("rateLimiterKeyResolver", RemoteAddrKeyResolver.BEAN_NAME);
        return keyResolverMap.getOrDefault(keyResolver, new PrincipalNameKeyResolver())
                .resolve(exchange);
    }
}
